# 功能说明：将PDF文件进行向量计算并持久化存储到向量数据库（chroma）
import os
import logging
from openai import OpenAI
import chromadb
import uuid
from utils import pdfSplitTest_Ch
from utils import pdfSplitTest_En
from typing import Optional, List

# 设置日志模版
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 配置类，统一管理所有模型配置
class ModelConfig:   
    def __init__(self, base_url: str, api_key: str, model: str):
        self.base_url = base_url
        self.api_key = api_key
        self.model = model

MODEL_CONFIGS = {
    # 国产大模型 OneAPI相关配置,通义千问为例
    "oneapi": ModelConfig(
        base_url="http://139.224.72.218:3000/v1",
        api_key="sk-GseYmJ8pX1D4I32323506e8fDf514a51A3C4B714FfD45aD9",
        model="text-embedding-v1"
    ),
    # 阿里通义千问大模型
    "qwen": ModelConfig(
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        api_key=os.getenv("DASHSCOPE_API_KEY"),
        model="text-embedding-v1"
    ),
    # 本地开源大模型 vLLM 方式
    # 本地开源大模型 xinference 方式
    # 本地开源大模型 Ollama 方式,bge-m3为例
    "ollama": ModelConfig(
        base_url="http://localhost:11434/v1",
        api_key="ollama",
        model="bge-m3:latest"
    ),
    # GPT大模型 OpenAI相关配置
    "openai": ModelConfig(
        base_url=os.getenv("OPENAI_API_KEY"),
        api_key=os.getenv("OPENAI_BASE_URL"),
        model="text-embedding-3-small"
    ),
}

# openai:调用gpt模型, qwen:调用阿里通义千问大模型, oneapi:调用oneapi方案支持的模型, ollama:调用本地开源大模型
LLM_TYPE = "qwen"

# 设置测试文本类型 Chinese 或 English
TEXT_LANGUAGE = 'Chinese'
INPUT_PDF = "input/健康档案.pdf"
# TEXT_LANGUAGE = 'English'
# INPUT_PDF = "input/deepseek-v3-1-4.pdf"

# 指定文件中待处理的页码，全部页码则填None
PAGE_NUMBERS=None
# PAGE_NUMBERS=[2, 3]

# 指定向量数据库chromaDB的存储位置和集合 根据自己的实际情况进行调整
CHROMADB_DIRECTORY = "chromaDB"  # chromaDB向量数据库的持久化路径
CHROMADB_COLLECTION_NAME = "demo001"  # 待查询的chromaDB向量数据库的集合名称


# get_embeddings方法计算向量
def get_embeddings(texts: List[str]) -> List[List[float]]:
    """
    获取文本的向量嵌入
    
    Args:
        texts: 待编码的文本列表
        
    Returns:
        向量嵌入列表，出错时返回空列表
    """

    # 字典自带get方法，如果key不存在则返回默认值
    config = MODEL_CONFIGS.get(LLM_TYPE, MODEL_CONFIGS['openai'])
    
    # 验证必要的配置
    if not config.api_key or not config.base_url:
        logger.error(f"{LLM_TYPE}配置不完整: API_KEY或BASE_URL缺失")
        return []
    
    try:
        client = OpenAI(base_url=config.base_url, api_key=config.api_key)
        response = client.embeddings.create(input=texts, model=config.model)
        return [item.embedding for item in response.data]
        
    except Exception as e:
        logger.error(f"生成向量时出错 ({LLM_TYPE}): {e}")
        return []


# 对文本按批次进行向量计算（分批计算提高效率）
def generate_vectors(data, max_batch_size=25):
    results = []
    for i in range(0, len(data), max_batch_size):#for循环，步长为max_batch_size
        batch = data[i:i + max_batch_size]
        # 调用向量生成get_embeddings方法  根据调用的API不同进行选择
        response = get_embeddings(batch)
        results.extend(response)
    return results


# 封装向量数据库chromadb类，提供两种方法
class MyVectorDBConnector:
    def __init__(self, collection_name, embedding_fn):
        # 申明使用全局变量
        global CHROMADB_DIRECTORY
        # 实例化一个chromadb对象
        # 设置一个文件夹进行向量数据库的持久化存储  路径为当前文件夹下chromaDB文件夹
        chroma_client = chromadb.PersistentClient(path=CHROMADB_DIRECTORY)
        # 创建一个collection数据集合
        # get_or_create_collection()获取一个现有的向量集合，如果该集合不存在，则创建一个新的集合
        self.collection = chroma_client.get_or_create_collection(
            name=collection_name)
        # embedding处理函数
        self.embedding_fn = embedding_fn

    # 添加文档到集合
    # 文档通常包括文本数据和其对应的向量表示，这些向量可以用于后续的搜索和相似度计算
    def add_documents(self, documents):
        self.collection.add(
            embeddings=self.embedding_fn(documents),  # 调用函数计算出文档中文本数据对应的向量
            documents=documents,  # 文档的文本数据
            ids=[str(uuid.uuid4()) for i in range(len(documents))]  # 文档的唯一标识符 自动生成uuid,128位  
        )
        
    # 检索向量数据库，返回包含查询结果的对象或列表，这些结果包括最相似的向量及其相关信息
    # query：查询文本
    # top_n：返回与查询向量最相似的前 n 个向量
    def search(self, query, top_n):
        try:
            results = self.collection.query(
                # 计算查询文本的向量，然后将查询文本生成的向量在向量数据库中进行相似度检索
                query_embeddings=self.embedding_fn([query]),
                n_results=top_n
            )
            return results
        except Exception as e:
            logger.info(f"检索向量数据库时出错: {e}")
            return []


# 封装文本预处理及灌库方法, 提供外部调用
def vectorStoreSave():
    global TEXT_LANGUAGE, CHROMADB_COLLECTION_NAME, INPUT_PDF, PAGE_NUMBERS

    # 根据语言选择对应的处理模块和查询语句
    language_config = {
        'Chinese': {
            'processor': pdfSplitTest_Ch,
            'query': "张三九的基本信息是什么"
        },
        'English': {
            'processor': pdfSplitTest_En,
            'query': "deepseek V3有多少参数"
        }
    }
     # 检查语言是否支持
    if TEXT_LANGUAGE not in language_config:
        logger.error(f"不支持的语言: {TEXT_LANGUAGE}")
        return
    
    # 获取对应语言的配置
    config = language_config[TEXT_LANGUAGE]
    
    # 1、获取处理后的文本数据
    # 演示测试对指定的全部页进行处理，其返回值为划分为段落的文本列表
    paragraphs = config['processor'].getParagraphs(
        filename=INPUT_PDF,
        page_numbers=PAGE_NUMBERS,
        min_line_length=1
    )
    
    # 2、将文本片段灌入向量数据库
    # 实例化一个向量数据库对象
    # 其中，传参collection_name为集合名称, embedding_fn为向量处理函数
    vector_db = MyVectorDBConnector(CHROMADB_COLLECTION_NAME, generate_vectors)
    # 向向量数据库中添加文档（文本数据、文本数据对应的向量数据）
    vector_db.add_documents(paragraphs)
    
    # 3、封装检索接口进行检索测试
    user_query = config['query']
    # 将检索出的5个近似的结果
    search_results = vector_db.search(user_query, 5)
    logger.info(f"检索向量数据库的结果: {search_results}")


if __name__ == "__main__":
    # 测试文本预处理及灌库
    vectorStoreSave()

